'''
Change annotations of object detection as the following (name it as 2007_xx.txt and place it in the root library in Aquarius, the same level with train.py).
eg.:
    need like: E:/Big_Datasets/water_surface/all-1114/all/VOCdevkit/VOC2007/JPEGImages/1664091257.87023.jpg 1131,430,1152,473,0 920,425,937,451,0
    now have:
        WaterScenes/detection/yolo/00001.txt(not use): 5 0.6190104166666667 0.525462962962963 0.033854166666666664 0.028703703703703703
        WaterScenes/detection/xmk/00001.xml

用于生成2007_train.txt/2007_val.txt文件，格式为：image_path(绝对路径) xmin,ymin,xmax,ymax,class_ind xmin,ymin,xmax,ymax,class_ind
仅用于WaterScenes数据集与Aquarius及其衍生算法
'''

import os
from tqdm import tqdm
import xml.etree.ElementTree as ET

def parse_waterscenes_annotation(data_path, file_type, anno_path, class_dict, use_difficult_bbox=True):
    """
    phase pascal flow annotation, eg:[image_global_path xmin,ymin,xmax,ymax,cls_id]
    :param data_path: eg: WaterScenes/
    :param file_type: eg: 'train''val''test'
    :param anno_path: path to ann file
    :class_dict: dict of class
    :param use_difficult_bbox: whither use different sample
    :return: batch size of data set
    """

    with open(os.path.join(data_path, file_type+'.txt'), 'r') as f:
        image_files = [os.path.join(data_path, x.strip()) for x in f.readlines()] # data_path/image/39484.jpg

    with open(anno_path, 'a') as f:
        for image_file in tqdm(image_files):
            annotation = os.path.abspath(image_file)
            label_path = os.path.join(data_path, 'detection/xml/', os.path.basename(image_file)).replace('.jpg', '.xml')
            root = ET.parse(label_path).getroot()
            objects = root.findall('object')
            for obj in objects:
                difficult = obj.find("difficult").text.strip()
                if (not use_difficult_bbox) and (int(difficult) == 1): # difficult 表示是否容易识别，0表示容易，1表示困难
                    continue
                name = obj.find('name').text.strip()
                class_id = class_dict[name] if name in class_dict else print('class error:', label_path, name)
                bbox = obj.find('bndbox')
                xmin = bbox.find('xmin').text.strip()
                ymin = bbox.find('ymin').text.strip()
                xmax = bbox.find('xmax').text.strip()
                ymax = bbox.find('ymax').text.strip()
                annotation += ' ' + ','.join([xmin, ymin, xmax, ymax, str(class_id)])
            annotation += '\n'
            # print(annotation)
            f.write(annotation)
    return len(image_files)

if __name__ =="__main__":
    DATA_PATH = "/mnt/g/WaterScenes/" # WSL
    CLASSES = "./model_data/waterscenes_benchmark.txt"

    with open(CLASSES, 'r') as f:
        class_list = [x.strip() for x in f.readlines()]

    class_dict = {}
    for i, class_ in enumerate(class_list):
        class_dict[class_] = i
    print(class_dict)

    # train_set
    train_annotation_path = '2007_train.txt'
    if os.path.exists(train_annotation_path):
        os.remove(train_annotation_path)

    # val_set
    val_annotation_path = '2007_val.txt'
    if os.path.exists(val_annotation_path):
        os.remove(val_annotation_path)

    # # test_set
    # test_annotation_path = '2007_test.txt'
    # if os.path.exists(test_annotation_path):
    #     os.remove(test_annotation_path)

    len_train = parse_waterscenes_annotation(DATA_PATH, "train", train_annotation_path, class_dict) # 从train.txt提取注释
    len_val = parse_waterscenes_annotation(DATA_PATH, "val", val_annotation_path, class_dict) # val.txt提取注释
    # len_test = parse_waterscenes_annotation(DATA_PATH, "test", test_annotation_path, class_dict) # val.test提取注释

    print("The number of images for train and test are :train : {0} | val : {1}".format(len_train, len_val))
